using System;
using System.Collections.Generic;
using System.Text;
using System.Runtime.InteropServices;
using System.ComponentModel;
using Microsoft.Win32.SafeHandles;

namespace usbburn
{
  public class NtDll
  {
    public class NtStatusException : Exception
    {
      const uint FORMAT_MESSAGE_ALLOCATE_BUFFER = 0x00000100;
      const uint FORMAT_MESSAGE_IGNORE_INSERTS = 0x00000200;
      const uint FORMAT_MESSAGE_FROM_SYSTEM = 0x00001000;
      const uint FORMAT_MESSAGE_FROM_HMODULE = 0x00000800;

      const uint LANG_NEUTRAL = 0x00;
      const uint SUBLANG_DEFAULT = 0x01;

      [DllImport("kernel32.dll")]
      static extern IntPtr LoadLibrary(string lpFileName);

      // the version, the sample is built upon:
      [DllImport("Kernel32.dll", SetLastError = true)]
      static extern uint FormatMessage(uint dwFlags, IntPtr lpSource,
         int dwMessageId, uint dwLanguageId, ref IntPtr lpBuffer,
         uint nSize, IntPtr pArguments);

      [DllImport("kernel32.dll", SetLastError = true)]
      static extern IntPtr LocalFree(IntPtr hMem);

      [DllImport("kernel32.dll", SetLastError = true)]
      static extern bool FreeLibrary(IntPtr hModule);

      static uint MAKELANGID(uint p, uint s)
      {
        return (uint)((((UInt16)(s)) << 10) | (UInt16)(p));
      }

      public NtStatusException(int err)
        : base(DisplayError(err))
      {
      }

      static string DisplayError(int NTStatusMessage)
      {
        IntPtr lpMessageBuffer = IntPtr.Zero;
        IntPtr Hand = LoadLibrary("NTDLL.DLL");

        uint dwChars = FormatMessage(
          FORMAT_MESSAGE_ALLOCATE_BUFFER |
          FORMAT_MESSAGE_FROM_SYSTEM |
          FORMAT_MESSAGE_FROM_HMODULE,
          Hand,
          NTStatusMessage,
          MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
          ref lpMessageBuffer,
          0,
          IntPtr.Zero);

        if (dwChars == 0)
          Marshal.ThrowExceptionForHR(Marshal.GetLastWin32Error());

        string sRet = Marshal.PtrToStringAnsi(lpMessageBuffer);

        // Free the buffer allocated by the system.
        LocalFree(lpMessageBuffer);
        FreeLibrary(Hand);

        return sRet;
      }
    }

    public const uint FILE_SYNCHRONOUS_IO_NONALERT = 0x00000020;

    public const uint SYNCHRONIZE = 0x00100000;

    public const uint GENERIC_READ = 0x80000000;
    public const uint GENERIC_WRITE = 0x40000000;

    public static bool NT_SUCCESS(int Status)
    {
      return (Status) >= 0;
    }

    /// <summary>
    /// Don't forget call FreeObjectAttributes, CLR specific!!!
    /// </summary>
    /// <param name="oa"></param>
    /// <param name="n"></param>
    /// <param name="a"></param>
    /// <param name="r"></param>
    /// <param name="s"></param>
    public static void InitializeObjectAttributes(ref OBJECT_ATTRIBUTES oa, UNICODE_STRING n, UInt32 a, UInt32 r, IntPtr s)
    {
      oa.Length = Marshal.SizeOf(typeof(OBJECT_ATTRIBUTES));
      oa.RootDirectory = r;
      oa.ObjectName = Marshal.AllocHGlobal(Marshal.SizeOf(oa.ObjectName));
      Marshal.StructureToPtr(n, oa.ObjectName, false);
      oa.Attributes = a;
      oa.SecurityDescriptor = s;
      oa.SecurityQualityOfService = IntPtr.Zero;
    }

    public static void FreeObjectAttributes(ref OBJECT_ATTRIBUTES oa)
    {
      if(oa.ObjectName!=IntPtr.Zero)
        Marshal.FreeHGlobal(oa.ObjectName);
    }

    [DllImport("ntdll.dll",
    SetLastError = true,
    CharSet = CharSet.Unicode,
    CallingConvention = CallingConvention.Winapi)]
    public static extern int NtOpenFile(
        ref IntPtr FileHandle,
        uint DesiredAccess,
        ref OBJECT_ATTRIBUTES ObjectAttributes,
        ref IO_STATUS_BLOCK IoStatusBlock,
        uint ShareAccess,
        uint OpenOptions);

    [DllImport("ntdll.dll", CharSet = CharSet.Unicode)]
    public static extern void RtlInitUnicodeString(
      ref UNICODE_STRING DestinationString,
      string SourceString);

    [DllImport("ntdll.dll", CharSet = CharSet.Unicode)]
    public static extern int NtClose(
      IntPtr Handle
      );

    [DllImport("kernel32.dll", SetLastError=true)]
    public static extern bool ReadFile(IntPtr hFile, byte[] lpBuffer,
       int nNumberOfBytesToRead, ref int lpNumberOfBytesRead, IntPtr lpOverlapped);

    [DllImport("kernel32.dll", SetLastError = true)]
    public static extern Boolean WriteFile(IntPtr fFile, Byte[] lpBuffer, int nNumberOfBytesToWrite,
            ref int lpNumberOfBytesWritten, IntPtr lpOverlapped);

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
    public struct IO_STATUS_BLOCK
    {
      public uint Status;
      public uint Information;
    }

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
    public struct UNICODE_STRING
    {
      public ushort Length;
      public ushort MaximumLength;
      public string Buffer;
    }

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
    public struct OBJECT_ATTRIBUTES
    {
      public int Length;
      public uint RootDirectory;
      public IntPtr ObjectName;
      public uint Attributes;
      public IntPtr SecurityDescriptor;        // Points to type SECURITY_DESCRIPTOR
      public IntPtr SecurityQualityOfService;  // Points to type SECURITY_QUALITY_OF_SERVICE
    }
  }
}
